import socket
import select
import threading
import logging
import sys
import time

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class ProxyServer:
    def __init__(self, host='127.0.0.1', port=8080, buffer_size=8192):
        self.host = host
        self.port = port
        self.buffer_size = buffer_size
        self.server_socket = None
        self.running = False
        
    def create_server_socket(self):
        try:
            self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
            self.server_socket.settimeout(60)  # 60秒超时
            return True
        except Exception as e:
            logger.error(f"Failed to create server socket: {e}")
            return False
        
    def start(self):
        if not self.create_server_socket():
            return
        
        try:
            self.server_socket.bind((self.host, self.port))
            self.server_socket.listen(100)  # 增加队列长度
            self.running = True
            logger.info(f"Proxy server started on {self.host}:{self.port}")
            
            while self.running:
                try:
                    client_socket, addr = self.server_socket.accept()
                    logger.info(f"New connection from {addr}")
                    client_socket.settimeout(30)  # 30秒客户端超时
                    client_thread = threading.Thread(
                        target=self.handle_client,
                        args=(client_socket, addr),
                        daemon=True
                    )
                    client_thread.start()
                except socket.timeout:
                    continue
                except Exception as e:
                    logger.error(f"Error accepting connection: {e}")
                    time.sleep(1)  # 避免CPU过载
                    
        except Exception as e:
            logger.error(f"Server error: {e}")
        finally:
            self.stop()
                
    def handle_client(self, client_socket, addr):
        remote_socket = None
        try:
            data = client_socket.recv(self.buffer_size)
            if not data:
                logger.warning(f"No data received from {addr}")
                return
            
            try:
                # 这里替换成你想要连接的实际服务器地址和端口
                remote_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                remote_socket.settimeout(30)
                remote_socket.connect(('127.0.0.1', 1080))  # 修改为你的目标服务器
                remote_socket.send(data)
                
                while self.running:
                    readable, _, exceptional = select.select(
                        [client_socket, remote_socket],
                        [],
                        [client_socket, remote_socket],
                        1.0
                    )
                    
                    if exceptional:
                        logger.error(f"Socket exception for {addr}")
                        break
                        
                    for sock in readable:
                        other = remote_socket if sock is client_socket else client_socket
                        try:
                            data = sock.recv(self.buffer_size)
                            if not data:
                                return
                            other.send(data)
                        except socket.error as e:
                            logger.error(f"Socket error while forwarding: {e}")
                            return
                            
            except socket.error as e:
                logger.error(f"Remote connection error: {e}")
                return
                    
        except Exception as e:
            logger.error(f"Error handling client {addr}: {e}")
        finally:
            self.close_socket(client_socket)
            if remote_socket:
                self.close_socket(remote_socket)
            logger.info(f"Connection closed for {addr}")
    
    def close_socket(self, sock):
        try:
            sock.shutdown(socket.SHUT_RDWR)
        except:
            pass
        try:
            sock.close()
        except:
            pass
            
    def stop(self):
        self.running = False
        if self.server_socket:
            self.close_socket(self.server_socket)
        logger.info("Proxy server stopped")

def main():
    proxy = ProxyServer(port=8080)  # 修改为你想要的端口
    try:
        proxy.start()
    except KeyboardInterrupt:
        logger.info("Shutting down proxy server")
        proxy.stop()
    except Exception as e:
        logger.error(f"Unexpected error: {e}")
        proxy.stop()
        sys.exit(1)

if __name__ == '__main__':
    main() 